# Copyright 2020 - 2021 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

EXPECTED_ANSWERS = [
    {  # test answers for PyTorch 1.6
        "integration_classification_2d": {
            "losses": [0.776835828070428, 0.1615355300011149, 0.07492854832938523, 0.04591309238865877],
            "best_metric": 0.9999184380485994,
            "infer_prop": [1029, 896, 980, 1033, 961, 1046],
        },
        "integration_segmentation_3d": {
            "losses": [
                0.5367561340332031,
                0.478084459900856,
                0.4581540793180466,
                0.44623913466930387,
                0.42341493666172025,
                0.42569945752620697,
            ],
            "best_metric": 0.9295084029436111,
            "infer_metric": 0.9296411260962486,
            "output_sums": [
                0.14302121377204619,
                0.15321686701244813,
                0.15267064069005093,
                0.1408481434833016,
                0.18862719991649474,
                0.16992848513054068,
                0.1479306037291329,
                0.1691071594535633,
                0.15804366588267224,
                0.18019304183940157,
                0.1635089455927468,
                0.16851606024285842,
                0.1454348651039073,
                0.11584957890961554,
                0.16255468027312903,
                0.20118089432240313,
                0.176187783307603,
                0.1004243279488101,
                0.19385348502657657,
                0.2030768555124136,
                0.196251372926592,
                0.20823046240222043,
                0.1631389353339986,
                0.13299661219478043,
                0.14917081129077908,
                0.14383374638201593,
                0.23050183928776746,
                0.1614747942341212,
                0.14913436515470202,
                0.10443081170610946,
                0.11978674347415241,
                0.13126176432899028,
                0.11570832453348577,
                0.15306806147195887,
                0.163673089782912,
                0.19394971756732426,
                0.22197501007172804,
                0.1812147930033603,
                0.19051659118682873,
                0.0774867922747158,
            ],
        },
        "integration_workflows": {
            "best_metric": 0.9219646483659745,
            "infer_metric": 0.921751058101654,
            "output_sums": [
                0.14183664321899414,
                0.1513957977294922,
                0.13804054260253906,
                0.13356828689575195,
                0.18456125259399414,
                0.16363763809204102,
                0.14090299606323242,
                0.16649389266967773,
                0.15651893615722656,
                0.17655134201049805,
                0.16116666793823242,
                0.1644763946533203,
                0.14383649826049805,
                0.11055326461791992,
                0.16080379486083984,
                0.19629907608032227,
                0.17441415786743164,
                0.053577423095703125,
                0.19043779373168945,
                0.19904804229736328,
                0.19526052474975586,
                0.20304107666015625,
                0.16030025482177734,
                0.13170623779296875,
                0.15118932723999023,
                0.13686418533325195,
                0.22668886184692383,
                0.1611471176147461,
                0.1472463607788086,
                0.10427379608154297,
                0.11962461471557617,
                0.1305704116821289,
                0.11204910278320312,
                0.15171337127685547,
                0.15962505340576172,
                0.18976259231567383,
                0.21649456024169922,
                0.17761802673339844,
                0.18516874313354492,
                0.03636503219604492,
            ],
            "best_metric_2": 0.9219559609889985,
            "infer_metric_2": 0.9217371672391892,
            "output_sums_2": [
                0.14187288284301758,
                0.15140819549560547,
                0.13802719116210938,
                0.1335887908935547,
                0.18454980850219727,
                0.1636652946472168,
                0.14091157913208008,
                0.16653108596801758,
                0.15651702880859375,
                0.17658615112304688,
                0.1611957550048828,
                0.16448307037353516,
                0.14385128021240234,
                0.1105203628540039,
                0.16085100173950195,
                0.19626951217651367,
                0.17442035675048828,
                0.053586483001708984,
                0.19042730331420898,
                0.1990523338317871,
                0.1952815055847168,
                0.20303773880004883,
                0.16034317016601562,
                0.13172531127929688,
                0.15118741989135742,
                0.1368694305419922,
                0.22667837142944336,
                0.16119050979614258,
                0.14726591110229492,
                0.10426473617553711,
                0.11961841583251953,
                0.13054800033569336,
                0.11203193664550781,
                0.15172529220581055,
                0.15963029861450195,
                0.18975019454956055,
                0.21646499633789062,
                0.17763566970825195,
                0.18517112731933594,
                0.03638744354248047,
            ],
        },
    },
    {  # test answers for PyTorch 1.7
        "integration_classification_2d": {
            "losses": [0.777176220515731, 0.16019743723664315, 0.07480076164197011, 0.045643698364780966],
            "best_metric": 0.9999418774120775,
            "infer_prop": [1030, 897, 980, 1033, 960, 1048],
        },
        "integration_segmentation_3d": {
            "losses": [
                0.5427072256803512,
                0.46434969305992124,
                0.45358552038669586,
                0.4363856494426727,
                0.42080804109573366,
                0.42058534920215607,
            ],
            "best_metric": 0.9292903542518616,
            "infer_metric": 0.9306288316845894,
            "output_sums": [
                0.14192493409895743,
                0.15182314591386872,
                0.15143080738742032,
                0.13972497034181824,
                0.18790884439406313,
                0.16933812661492562,
                0.14664343345928132,
                0.1678599094806423,
                0.1568852615222309,
                0.17882538307200632,
                0.16226220644853354,
                0.16756325103417588,
                0.1449974856885373,
                0.1160602083671129,
                0.1614830941632057,
                0.20060717335382267,
                0.17543495742507476,
                0.10308107883493946,
                0.19289222718691168,
                0.20225689438356148,
                0.19587806881756237,
                0.20773073456322155,
                0.16193015294299506,
                0.13181961683097554,
                0.14850995284454005,
                0.14238637655756,
                0.2307113922277095,
                0.1608335768948913,
                0.1480752874532259,
                0.1038477413165911,
                0.11880665574424197,
                0.13084873656303445,
                0.1141965805147642,
                0.1531586543003841,
                0.16275008603701097,
                0.19320476187766733,
                0.2217811250932611,
                0.18027048819200148,
                0.18958803602663193,
                0.08653716931250294,
            ],
        },
        "integration_workflows": {
            "best_metric": 0.9217087924480438,
            "infer_metric": 0.9214379042387009,
            "output_sums": [
                0.14209461212158203,
                0.15126705169677734,
                0.13800382614135742,
                0.1338181495666504,
                0.1850571632385254,
                0.16372442245483398,
                0.14059066772460938,
                0.16674423217773438,
                0.15653657913208008,
                0.17690563201904297,
                0.16154909133911133,
                0.16521310806274414,
                0.14388608932495117,
                0.1103353500366211,
                0.1609959602355957,
                0.1967010498046875,
                0.1746964454650879,
                0.05329275131225586,
                0.19098854064941406,
                0.19976520538330078,
                0.19576644897460938,
                0.20346736907958984,
                0.1601848602294922,
                0.1316051483154297,
                0.1511220932006836,
                0.13670969009399414,
                0.2276287078857422,
                0.1611800193786621,
                0.14751672744750977,
                0.10413789749145508,
                0.11944007873535156,
                0.1305546760559082,
                0.11204719543457031,
                0.15145111083984375,
                0.16007614135742188,
                0.1904129981994629,
                0.21741962432861328,
                0.17812013626098633,
                0.18587207794189453,
                0.03605222702026367,
            ],
            "best_metric_2": 0.9210659921169281,
            "infer_metric_2": 0.9208109736442566,
            "output_sums_2": [
                0.14227628707885742,
                0.1515035629272461,
                0.13819408416748047,
                0.13402271270751953,
                0.18525266647338867,
                0.16388607025146484,
                0.14076614379882812,
                0.16694307327270508,
                0.15677356719970703,
                0.1771831512451172,
                0.16172313690185547,
                0.1653728485107422,
                0.14413118362426758,
                0.11057281494140625,
                0.16121912002563477,
                0.19680166244506836,
                0.1748638153076172,
                0.053426265716552734,
                0.19117307662963867,
                0.19996356964111328,
                0.1959366798400879,
                0.20363712310791016,
                0.16037797927856445,
                0.13180780410766602,
                0.1513657569885254,
                0.13686084747314453,
                0.2277364730834961,
                0.16137409210205078,
                0.1476879119873047,
                0.10438394546508789,
                0.11967992782592773,
                0.13080739974975586,
                0.11226606369018555,
                0.15168476104736328,
                0.1602616310119629,
                0.190582275390625,
                0.21756458282470703,
                0.17825984954833984,
                0.18604803085327148,
                0.036206722259521484,
            ],
        },
    },
]


def test_integration_value(test_name, key, data, rtol=1e-2):
    for expected in EXPECTED_ANSWERS:
        if test_name not in expected:
            continue
        value = expected[test_name][key]
        if np.allclose(data, value, rtol=rtol):
            return True
    raise ValueError(f"no matched results for {test_name}, {key}. {data}.")
